// 文件：src/ch08/code-tools.ts
import { Tool } from "@langchain/core/tools";
import { z } from "zod";

export class CodeExecutionTool extends Tool {
  name = "execute_code";
  description = "执行 Python 代码，支持数据分析、图表生成等";
  schema = z.object({
    code: z.string().describe("要执行的 Python 代码"),
    language: z.enum(["python"]).default("python").describe("编程语言")
  });

  private sandbox = new Map<string, any>();

  async _call(input: { code?: string; language?: string }) {
    const language = input.language ?? "python";
    if (language !== "python") {
      return "目前只支持 Python 代码执行";
    }

    if (!input.code || input.code.trim() === "") {
      return "执行错误：缺少要执行的代码";
    }

    try {
      // 简化的 Python 代码执行（实际应使用安全的沙箱）
      const result = await this.executePython(input.code);
      return `执行结果：\n${result}`;
    } catch (error) {
      return `执行错误：${error}`;
    }
  }

  private async executePython(code: string): Promise<string> {
    // 安全检查
    const dangerous = ['import os', 'import sys', 'eval(', 'exec(', '__import__'];
    for (const d of dangerous) {
      if (code.includes(d)) {
        throw new Error(`禁止使用：${d}`);
      }
    }

    // 改进的模拟执行：识别使用 pandas 读取 CSV 并生成统计摘要的代码
    if (code.includes('pd.read_csv') || code.includes('describe(')) {
      // 解析嵌入的 csv_text 三引号字符串
      const match = code.match(/csv_text\s*=\s*"""([\s\S]*?)"""/);
      if (match) {
        const csv = match[1].trim();
        // 兼容传入数据中使用了转义的 \n，而不是实际换行
        const normalized = csv.replace(/\\n/g, '\n');
        const lines = normalized.split(/\n/).filter(l => l.trim().length > 0);
        if (lines.length === 0) return "执行错误：CSV 内容为空";
        const headers = lines[0].split(',').map(h => h.trim());
        const rows = lines.slice(1).map(line => line.split(',').map(v => v.trim()));

        const colCount = headers.length;
        const rowCount = rows.length;
        const colNames = headers.join(',');

        // 计算简单统计
        const stats: string[] = [];
        for (let c = 0; c < colCount; c++) {
          const colName = headers[c];
          const values = rows.map(r => r[c]);
          const nums = values.map(v => Number(v)).filter(v => !isNaN(v));
          if (nums.length === values.length && nums.length > 0) {
            const count = nums.length;
            const sum = nums.reduce((a, b) => a + b, 0);
            const mean = sum / count;
            const min = Math.min(...nums);
            const max = Math.max(...nums);
            stats.push(`${colName}: count=${count}, mean=${Number(mean.toFixed(2))}, min=${min}, max=${max}`);
          } else {
            const count = values.length;
            const freq: Record<string, number> = {};
            for (const v of values) freq[v] = (freq[v] || 0) + 1;
            const unique = Object.keys(freq).length;
            const top = Object.entries(freq).sort((a, b) => b[1] - a[1])[0];
            const topStr = top ? `${top[0]} (${top[1]})` : '无';
            stats.push(`${colName}: count=${count}, unique=${unique}, top=${topStr}`);
          }
        }

        return [
          '数据概览：',
          `行数：${rowCount}`,
          `列数：${colCount}`,
          `列名： ${colNames}`,
          '',
          '统计摘要：',
          ...stats
        ].join('\n');
      }
      // 无法解析嵌入 CSV 文本时，返回一般性提示
      return "已识别数据分析代码，但未找到嵌入的 CSV 文本";
    }

    // 其他简单模拟
    if (code.includes('import matplotlib')) {
      return "matplotlib 已导入，可以生成图表";
    }
    if (code.includes('print(')) {
      return "模拟输出：代码执行成功";
    }
    return "代码执行完成";
  }
}

// 数据分析工具
export class DataAnalysisTool extends Tool {
  name = "analyze_data";
  description = "分析数据，生成统计信息和可视化";
  schema = z.object({
    data: z.string().describe("CSV 格式的数据或数据描述"),
    analysis_type: z.enum(["summary", "visualization", "correlation"]).describe("分析类型")
  });

  async _call(input: { data: string; analysis_type: string }) {
    const type = input.analysis_type ?? "summary";
    const code = this.generateAnalysisCode(input.data, type);
    const execTool = new CodeExecutionTool();
    return await execTool._call({ code, language: "python" });
  }

  private generateAnalysisCode(data: string, type: string): string {
    switch (type) {
      case "summary":
        return `
import pandas as pd
import io
csv_text = """${data}"""
df = pd.read_csv(io.StringIO(csv_text))
print("数据概览：")
print(f"行数：{len(df)}")
print(f"列数：{len(df.columns)}")
print("列名：", ",".join(df.columns))
print("\\n统计摘要：")
print(df.describe(include='all'))
`;
      case "visualization":
        return `
import matplotlib.pyplot as plt
import seaborn as sns
# 生成示例图表
plt.figure(figsize=(10, 6))
plt.title("数据可视化")
plt.show()
`;
      default:
        return "print('分析完成')";
    }
  }
}
